{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "503bb02c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from agents import Agent, WebSearchTool, trace, Runner, ItemHelpers, function_tool\n",
    "from agents.model_settings import ModelSettings\n",
    "from pydantic import BaseModel, Field\n",
    "from dotenv import load_dotenv\n",
    "import asyncio\n",
    "import sendgrid\n",
    "import os\n",
    "import json\n",
    "from sendgrid.helpers.mail import Mail, Email, To, Content\n",
    "from typing import Dict\n",
    "import gradio as gr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e937e176",
   "metadata": {},
   "outputs": [],
   "source": [
    "load_dotenv(override=True)\n",
    "\n",
    "MODEL = \"gpt-4o-mini\"\n",
    "HOW_MANY_SEARCHES = 2\n",
    "STEPS = {\n",
    "    \"clarification_agent\": \"💭 Understanding your research needs...\",\n",
    "    \"question_generation_agent\": \"🎯 Crafting smart search strategies...\",\n",
    "    \"search_agent\": \"🔍 Exploring the web for insights...\",\n",
    "    \"writer_agent\": \"✨ Weaving findings into a beautiful report...\",\n",
    "    \"email_agent\": \"📬 Delivering your report to your inbox...\"\n",
    "}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c74ed82b",
   "metadata": {},
   "outputs": [],
   "source": [
    "class WebSearchItem(BaseModel):\n",
    "    reason: str = Field(description=\"Your reasoning for why this search is important to the query.\")\n",
    "\n",
    "    query: str = Field(description=\"The search term to use for the web search.\")\n",
    "\n",
    "\n",
    "class WebSearchPlan(BaseModel):\n",
    "    searches: list[WebSearchItem] = Field(description=\"A list of web searches to perform to best answer the query.\")\n",
    "    \n",
    "\n",
    "class ReportData(BaseModel):\n",
    "    short_summary: str = Field(description=\"A short 2-3 sentence summary of the findings.\")\n",
    "\n",
    "    markdown_report: str = Field(description=\"The final report\")\n",
    "\n",
    "    follow_up_questions: list[str] = Field(description=\"Suggested topics to research further\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7312f774",
   "metadata": {},
   "outputs": [],
   "source": [
    "CLARIFICATION_INSTRUCTIONS = \"\"\"\n",
    "You are a helpful research assistant specializing in understanding user needs.\n",
    "\n",
    "Analyze the user's research topic and generate 2-4 clarifying questions that will help you understand what they really need from this research.\n",
    "\n",
    "Your questions must be directly relevant to the specific topic they mentioned. Think about what would be ambiguous or unclear about that particular subject and ask about those aspects.\n",
    "\n",
    "Be concise and friendly.\n",
    "\"\"\"\n",
    "\n",
    "clarification_agent = Agent(\n",
    "    name=\"Clarification agent\",\n",
    "    instructions=CLARIFICATION_INSTRUCTIONS,\n",
    "    model=MODEL\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bff2a91a",
   "metadata": {},
   "outputs": [],
   "source": [
    "QUESTION_GENERATION_INSTRUCTIONS = f\"\"\"\n",
    "You are an expert research planner who generates comprehensive search queries.\n",
    "\n",
    "Based on the user's original topic and their clarifications, generate exactly {HOW_MANY_SEARCHES} diverse search queries.\n",
    "\n",
    "Your queries should:\n",
    "- Cover different aspects of the topic\n",
    "- Incorporate the specific angles the user mentioned\n",
    "- Range from broad overview to specific details\n",
    "- Consider different perspectives or approaches\n",
    "- Complement each other without excessive overlap\n",
    "\n",
    "For each query, provide a clear reason why that search is important and ensure the query itself is specific and searchable.\n",
    "\"\"\"\n",
    "\n",
    "question_generation_agent = Agent(\n",
    "    name=\"Question generation agent\",\n",
    "    instructions=QUESTION_GENERATION_INSTRUCTIONS,\n",
    "    output_type=WebSearchPlan,\n",
    "    model=MODEL\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d67a5b1b",
   "metadata": {},
   "outputs": [],
   "source": [
    "SEARCH_INSTRUCTIONS = \"\"\"\n",
    "You are a research assistant specialized in web search and summarization.\n",
    "\n",
    "You will receive a WebSearchPlan containing multiple search queries. For each query, use your WebSearchTool to search the web and produce a concise summary of the results.\n",
    "\n",
    "Requirements for each summary:\n",
    "- 2-3 paragraphs, less than 300 words\n",
    "- Capture only the main points, write succinctly\n",
    "- Focus on essence, ignore fluff\n",
    "- Information-dense and focused\n",
    "- No commentary, only factual summary\n",
    "\n",
    "Process all queries in the plan. Your summaries will be used to create a comprehensive report.\n",
    "\"\"\"\n",
    "\n",
    "search_agent = Agent(\n",
    "    name=\"Search agent\",\n",
    "    instructions=SEARCH_INSTRUCTIONS,\n",
    "    tools=[WebSearchTool(search_context_size=\"low\")],\n",
    "    model=MODEL,\n",
    "    model_settings=ModelSettings(tool_choice=\"required\")\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96d8edde",
   "metadata": {},
   "outputs": [],
   "source": [
    "WRITER_INSTRUCTIONS = \"\"\"\n",
    "You are a senior research analyst specializing in synthesizing information from multiple sources.\n",
    "\n",
    "You will receive the original research topic and multiple search summaries. Your task is to synthesize all information into one comprehensive report.\n",
    "\n",
    "Synthesize means: identify common themes, remove redundancies, organize by topic (not by source), create a flowing narrative with analysis and insights.\n",
    "\n",
    "Report requirements:\n",
    "- Length: 1000+ words\n",
    "- Format: Professional markdown with proper headings (##, ###)\n",
    "- Structure: Introduction, multiple body sections, conclusion\n",
    "- Style: Clear, authoritative, well-organized\n",
    "\n",
    "Return a ReportData object with:\n",
    "- short_summary: 2-3 sentences capturing key findings\n",
    "- markdown_report: The full detailed report in markdown\n",
    "- follow_up_questions: 3-5 suggested topics for deeper research\n",
    "\"\"\"\n",
    "\n",
    "writer_agent = Agent(\n",
    "    name=\"WriterAgent\",\n",
    "    instructions=WRITER_INSTRUCTIONS,\n",
    "    model=\"gpt-4o-mini\",\n",
    "    output_type=ReportData,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e902a8a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "EMAIL_INSTRUCTIONS = \"\"\"\n",
    "You are an email specialist agent receiving a ReportData object and user email address via handoff.\n",
    "\n",
    "The ReportData contains: short_summary, markdown_report, and follow_up_questions.\n",
    "The handoff message includes the user's email address.\n",
    "\n",
    "Your task:\n",
    "1. Extract the user's email address from the handoff context\n",
    "2. Convert the markdown_report to well-formatted professional HTML\n",
    "3. Create an engaging subject line based on the research topic (e.g., \"Your Research Report: [Topic]\")\n",
    "4. Use your send_email tool once to send the email to the user's address\n",
    "\n",
    "Execute immediately with the data provided.\n",
    "\"\"\"\n",
    "\n",
    "@function_tool\n",
    "def send_email(subject: str, html_body: str, to_email:str) -> Dict[str, str]:\n",
    "    \"\"\" \n",
    "    Send out an email with the given subject and HTML body.\n",
    "    \"\"\"\n",
    "    print(f\"Sending email to {to_email} with subject {subject}\", flush=True)\n",
    "    sg = sendgrid.SendGridAPIClient(api_key=os.environ.get('SENDGRID_API_KEY'))\n",
    "    from_email = Email(\"ed@edwarddonner.com\") # Change this to your verified email\n",
    "    to_email = To(to_email) # Change this to your email\n",
    "    content = Content(\"text/html\", html_body)\n",
    "    mail = Mail(from_email, to_email, subject, content).get()\n",
    "    sg.client.mail.send.post(request_body=mail)\n",
    "    return \"Email sent successfully to the user\"\n",
    "\n",
    "email_agent = Agent(\n",
    "    name=\"Email Agent\",\n",
    "    instructions=EMAIL_INSTRUCTIONS,\n",
    "    tools=[send_email],\n",
    "    model=\"gpt-4o-mini\",\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df7885dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "PLANNER_INSTRUCTIONS = \"\"\"\n",
    "You are the Deep Research Coordinator, a friendly AI assistant that helps users conduct comprehensive research.\n",
    "\n",
    "FIRST: Determine the user's intent:\n",
    "\n",
    "A. CASUAL CONVERSATION (greetings, questions about you, small talk):\n",
    "   - Respond naturally and warmly without using any tools\n",
    "   - Examples: \"Hi\", \"Hello\", \"How are you?\", \"What can you do?\", \"Who are you?\"\n",
    "   - Introduce yourself as a research assistant that helps conduct deep web research\n",
    "   - Invite them to ask you to research any topic they're curious about\n",
    "\n",
    "B. RESEARCH REQUEST (explicit or implicit request to research a topic):\n",
    "   - Keywords: \"research\", \"investigate\", \"find out about\", \"tell me about\", \"I want to know\", \"information on\", etc.\n",
    "   - Proceed with the research workflow below\n",
    "\n",
    "RESEARCH WORKFLOW:\n",
    "\n",
    "Check conversation history:\n",
    "- If the LAST ASSISTANT MESSAGE contains clarifying questions, the current user message is likely a RESPONSE. SKIP STEP 1, go to STEP 2.\n",
    "- Otherwise, START AT STEP 1.\n",
    "\n",
    "Execute these steps sequentially:\n",
    "\n",
    "1. CLARIFICATION (Skip if user is responding to previous clarifying questions): Use the clarification agent to ask the user 2-4 specific questions about their research topic. Wait for their response before proceeding.\n",
    "\n",
    "2. GENERATE SEARCH PLAN: Use the question generation agent with the topic and clarifications to create diverse search queries.\n",
    "\n",
    "3. EXECUTE SEARCHES: Use the search agent to perform all searches in the plan. Collect all summaries before proceeding.\n",
    "\n",
    "4. SYNTHESIZE REPORT: Use the writer agent with the topic and all search summaries to create a comprehensive report. It returns a ReportData object.\n",
    "\n",
    "5. AFTER WRITER AGENT - CRITICAL DECISION POINT:\n",
    "   ⚠️ MANDATORY: After writer_agent completes, you MUST follow these EXACT rules. NO EXCEPTIONS.\n",
    "   \n",
    "   First, check conversation history for an email address (format: xxx@xxx.xxx).\n",
    "   \n",
    "   OPTION A - Email found in history:\n",
    "   ⛔ DO NOT write any message\n",
    "   ⛔ DO NOT explain what you're doing\n",
    "   ⛔ DO NOT say anything about the report\n",
    "   ✅ IMMEDIATELY handoff to email agent (step 6)\n",
    "   \n",
    "   OPTION B - NO email in history:\n",
    "   ⛔ DO NOT explain the report\n",
    "   ⛔ DO NOT summarize findings\n",
    "   ⛔ DO NOT add greetings or conclusions\n",
    "   ✅ Write ONLY ONE sentence: \"Great! Your research report is ready. What email address should I send the full report to?\"\n",
    "   ✅ Then STOP and wait for user response\n",
    "   ✅ After receiving email, go to step 6\n",
    "   \n",
    "   FORBIDDEN after writer_agent: Summaries, explanations, \"great news\", \"I've completed\", or any commentary about the research.\n",
    "\n",
    "6. HANDOFF: Hand off to the Email Agent with the ReportData AND the user's email address. Include the email in your handoff message like: \"Please send the report to [user's email]\". Your job is complete after handoff.\n",
    "\n",
    "Wait for user input when needed. Ensure all searches complete before synthesis. Never skip the email collection step.\n",
    "\"\"\"\n",
    "\n",
    "planner_agent = Agent(\n",
    "    name=\"DeepResearchCoordinator\",\n",
    "    instructions=PLANNER_INSTRUCTIONS,\n",
    "    tools=[\n",
    "       clarification_agent.as_tool(\n",
    "           tool_name=\"clarification_agent\",\n",
    "           tool_description=\"Ask the user clarifying questions about their research topic to understand their needs better\"\n",
    "       ),\n",
    "       question_generation_agent.as_tool(\n",
    "           tool_name=\"question_generation_agent\",\n",
    "           tool_description=\"Generate diverse search queries based on the topic and user clarifications\"\n",
    "       ),\n",
    "        search_agent.as_tool(\n",
    "            tool_name=\"search_agent\",\n",
    "            tool_description=\"Execute web searches for the given queries and return summaries\"\n",
    "        ),\n",
    "        writer_agent.as_tool(\n",
    "            tool_name=\"writer_agent\",\n",
    "            tool_description=\"Synthesize all search results into a comprehensive research report\"\n",
    "        )\n",
    "    ],\n",
    "    handoffs=[email_agent],\n",
    "    handoff_description=\"Hand off to Email Agent to send the final research report\",\n",
    "    model=MODEL\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "214bfd3a",
   "metadata": {},
   "outputs": [],
   "source": [
    "async def deep_research_chat(message, history):\n",
    "    \"\"\"\n",
    "    Handle the deep research conversation with the planner agent.\n",
    "    Maintains conversation history for multi-turn interactions.\n",
    "    \"\"\"\n",
    "    try:      \n",
    "        yield \"✨ Thinking...\"\n",
    "        \n",
    "        with trace(\"DeepResearch\"):\n",
    "            history = [{\"role\": msg[\"role\"], \"content\": msg[\"content\"]} for msg in history]\n",
    "            messages = [*history, {\"role\": \"user\", \"content\": message}]\n",
    "            result = Runner.run_streamed(planner_agent, input=messages)\n",
    "            last_called_tool = None\n",
    "            report_shown = False\n",
    "            report = None\n",
    "\n",
    "            async for event in result.stream_events():\n",
    "                if event.type == \"raw_response_event\":\n",
    "                    continue\n",
    "                elif event.type == \"agent_updated_stream_event\":\n",
    "                    print(f\"Agent updated: {event.new_agent.name}\")\n",
    "                    continue\n",
    "                elif event.type == \"run_item_stream_event\":\n",
    "                    if event.item.type == \"tool_call_item\":\n",
    "                        print(f\"-- Tool was called {event.item.raw_item.name}\")\n",
    "                        last_called_tool = event.item.raw_item.name\n",
    "                        yield f\"{STEPS[last_called_tool]}\"\n",
    "                        await asyncio.sleep(1)\n",
    "                    elif event.item.type == \"tool_call_output_item\":\n",
    "                        print(f\"-- Tool output: {event.item.output}\")\n",
    "                        if last_called_tool == \"writer_agent\":\n",
    "                            result = json.loads(event.item.output)\n",
    "                            \n",
    "                            follow_up_list = \"\"\n",
    "                            for i, question in enumerate(result[\"follow_up_questions\"], 1):\n",
    "                                follow_up_list += f\"\\n{i}. {question}\"\n",
    "                            \n",
    "                            report =  f\"## 📊 Research Complete!\\n### Summary\\n{result[\"short_summary\"]}\\n### 📧 Email Status\\n✅ Full report will be sent to your email.\\n### 🔍 Suggested Follow-up Topics\\n{follow_up_list}\"\n",
    "                            yield report\n",
    "                            report_shown = True\n",
    "                    elif event.item.type == \"message_output_item\":\n",
    "                        if not report_shown:\n",
    "                            yield ItemHelpers.text_message_output(event.item)\n",
    "                        else:\n",
    "                            yield report + \"\\n\\n\" + ItemHelpers.text_message_output(event.item)\n",
    "            \n",
    "                             \n",
    "                \n",
    "    except Exception as e:\n",
    "        yield f\"❌ Error: {str(e)}\\n\\nPlease try again or check your configuration.\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4629c291",
   "metadata": {},
   "outputs": [],
   "source": [
    "with gr.Blocks(theme=gr.themes.Soft(), title=\"Deep Research Assistant\") as demo:\n",
    "    \n",
    "    gr.Markdown(\"\"\"\n",
    "    # 🔍 Deep Research Assistant\n",
    "    \n",
    "    An intelligent multi-agent system that performs comprehensive web research and engages in natural conversation.\n",
    "    \"\"\")\n",
    "    \n",
    "    with gr.Row():\n",
    "        with gr.Column(scale=2):\n",
    "            chatbot = gr.ChatInterface(\n",
    "                deep_research_chat,\n",
    "                type=\"messages\",\n",
    "                chatbot=gr.Chatbot(height=650, type=\"messages\"),\n",
    "                examples=[  \n",
    "                    \"Hello! What can you do?\",\n",
    "                    \"I want to research quantum computing in drug discovery\",\n",
    "                    \"Research the latest developments in renewable energy\",\n",
    "                    \"Tell me about AI applications in healthcare\"\n",
    "                ],\n",
    "                cache_examples=False,\n",
    "            )\n",
    "        \n",
    "        with gr.Column(scale=1):\n",
    "            gr.Markdown(\"\"\"\n",
    "            ### 💬 Chat Naturally\n",
    "            \n",
    "            You can **chat with me just like ChatGPT!**\n",
    "            \n",
    "            - Ask casual questions\n",
    "            - Get research topic suggestions\n",
    "            - Request deep research on any topic\n",
    "            - I'll adapt to what you need\n",
    "            \n",
    "            ---\n",
    "            \n",
    "            ### 🔬 Research Workflow\n",
    "            \n",
    "            When you ask me to research:\n",
    "            \n",
    "            1. 💭 **Clarification** - I'll ask questions to understand your needs\n",
    "            2. 🎯 **Smart Queries** - I'll generate targeted search strategies\n",
    "            3. 🔍 **Deep Search** - I'll explore multiple sources\n",
    "            4. ✨ **Synthesis** - I'll weave findings into a comprehensive report\n",
    "            5. 📧 **Email Collection** - I'll ask for your email address\n",
    "            6. 📬 **Delivery** - Full report sent to your inbox\n",
    "            \n",
    "            ---\n",
    "            \n",
    "            ### 💡 Tips\n",
    "            \n",
    "            **For Research:**\n",
    "            - Be specific about your topic\n",
    "            - Answer my clarifying questions\n",
    "            - Process takes 2-5 minutes\n",
    "            - I'll ask for your email to send the report\n",
    "            - Full report delivered via email\n",
    "            \n",
    "            **For Topic Ideas:**\n",
    "            - Tell me what you're working on\n",
    "            - Ask: \"Suggest research topics for [your project]\"\n",
    "            - I'll provide relevant suggestions\n",
    "            \n",
    "            ---\n",
    "            \n",
    "            ### 🚀 Try It Out\n",
    "            \n",
    "            - \"Hi, what can you help me with?\"\n",
    "            - \"Suggest topics for my AI project\"\n",
    "            - \"Research AI ethics in healthcare\"\n",
    "            \"\"\")\n",
    "\n",
    "# Launch the interface\n",
    "demo.launch()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
